__doc__ =   """ Functions and classes for crystallographic needs.
                This module requires the cctbx library.
            """
#Python
from operator import xor, not_
from copy import deepcopy
from itertools import izip
from collections import defaultdict
from math import pow, sqrt, pi
from re import compile

#External
from cctbx import sgtbx, crystal
from numpy import float_, array, linalg, dot, squeeze,\
                  mean, c_, ones, zeros, r_, vstack

#ZenPDB
from Einput import einput
from Entities import Structure
from MultiEntity import MultiEntity
from Crunchers import kdtree, asa_loop, cnt_loop
from Math import sphere_points

def get_symmetry(cryst1, space_group):
        unit_cell = float_(cryst1.split())        # removes before and after spaces too.
        sg_string = space_group.replace('H', 'R') # phenix requirement
        symmetry = crystal.symmetry(unit_cell =unit_cell, space_group_symbol =sg_string) # call phenix
        uc = symmetry.unit_cell()                           # get unit cell
        fmx = array(uc.fractionalization_matrix()).reshape((3,3))   # get fractional/orthogonal conversion
        omx = array(uc.orthogonalization_matrix()).reshape((3,3))   # matrices as 3x3 scale matrix
        valid_sg = sgtbx.space_group_symbols(sg_string, '').hall()  # convert to hall
        parsed_sg = sgtbx.parse_string(valid_sg)
        sg = sgtbx.space_group(parsed_sg)                           # space group
        mxs = []
        for mx in sg:                                     # for matrix in symmetry matrices
            mxr = mx.as_4x4_rational().as_float().as_list_of_lists()  # 4x4 = 3x3 + 3x1 + 0 0 0 1
            mxs.append(mxr)                                           # the bottom row 0001 for inverse/dot
        mxs = array(mxs)
        cmx = []

def axis_type(ents, level ='A', xtra_key =None, cnt_key =None):
    """ Determines the parameters of a crystallographic axis, its: fold, sense and screw.
    """
    entities = einput(ents, level)          # the level is arbitrary but done at the atom
    xtra_key = (xtra_key or 'SYM_ACT_X')    # level elimnates the need to transfer
    cnt_key =  (cnt_key  or 'CNT_ACT_X')    # contacts on residues.
    for entity in entities:
        try:
            sg = entity.get_parent('S').header['space_group']
        except KeyError:
            sg = None
        symmetry = None
        symmetries = []
        contacts = entity.xtra.get(cnt_key)
        if contacts: # skip get-None
            for (full_id, (dist, symop)) in contacts.iteritems():
                if not symop.strip():    # ACT contact inside the ASU
                    (fold, sense, screw) = (1, 0, 0)
                else:
                    try:                            # try to get the symmetry number
                        symop = int(symop[1:5]) - 1 # python and cctbx count symops from 0
                    except ValueError:
                        pass                        # we got ncont
                    (fold, sense, screw) = get_axis_type(symop, sg)
                symmetry = "%s:%s:%s" % (fold, sense, screw)
                symmetries.append(symmetry)
            entity.xtra[xtra_key] = symmetries

def get_axis_type(symop_string, sg_string =None, convention =''):
    """ given given symop as xyz or a symmetry number for the appropriate space group
        returns the type and sense of the axis.
    """
    if sg_string and isinstance(symop_string, int):   # try symop # from sg first
        sg_string = sg_string.replace('H', 'R')
        valid_sg = sgtbx.space_group_symbols(sg_string, convention).hall()  # convert to hall
        parsed_sg = sgtbx.parse_string(valid_sg)
        sg = sgtbx.space_group(parsed_sg)                                   # get SG
        rot_mx = sg[symop_string]                                           # get rotation matrix
    else:           # no sg parse symop
        parsed_symop = sgtbx.parse_string(symop_string)
        rot_mx = sgtbx.rt_mx(parsed_symop)
    r_info = sgtbx.rot_mx_info(rot_mx.r())          # rotation info
    t_info = sgtbx.translation_part_info(rot_mx)    # translation info
    fold = r_info.type()
    sense = r_info.sense()
    screw = int(not_((t_info.intrinsic_part().is_zero())))  # am I screw?
    return (fold, sense, screw)

#EOF